#include "rockchip/rkmodel_manager.h"
#include "base_util/utils.h"
#include "base_util/math_compute.h"
#include "stream/osd_text.h"
#include "common/face_tool.h"
#include "infer_demo.h"
using namespace httplib;


int rec_car(cv::Mat& img, std::vector<cv::Rect>& boxs) {
  int ret = 0;

  // 识别车辆
  std::string b64_encode;
  cv_image_encode(img, b64_encode);
  Json::Value data;
  data["b64_img"] = b64_encode;
  std::string ret_body, msg = data.toStyledString();

  std::string httpServer = "http://2.2.2.81:9966";
  std::string uri = "/api/rec2";
  post_data_base(ret_body, msg, httpServer, uri, -1, 1);
  Json::Value ret_data;
  ret = parse_json(ret_body, ret_data, 0);
  if (ret != 0 || ret_data.get("code", Json::Value(-1)).asInt() != 0) {
    std::cout << "parse_json:" << ret << std::endl;
    return ret;
  }

  for (auto& item : ret_data["items"]) {
    int x = item["x"].asInt();
    int y = item["y"].asInt();
    int w = item["w"].asInt();
    int h = item["h"].asInt();
    if (x < 0) {x = 0;}
    if (y < 0) {y = 0;}
    if (x + w > img.cols) {w = img.cols - x;}
    if (y + h > img.rows) {h = img.rows - y;}
    boxs.push_back(cv::Rect(x,y,w,h));
  }
  return 0;
}



int infer_det_mp4_plate(ai::RKModelManager* model, stream::CVOsdText* osd, std::string mp4_path) {
  int ret = 0;
  cv::VideoCapture reader = cv::VideoCapture(mp4_path);
  if (!reader.isOpened()) { 
    printf("ERROR. infer_mp4 !reader->isOpened() mp4_path: %s\n", mp4_path);
    return -1;
  }
  
  int width = reader.get(cv::CAP_PROP_FRAME_WIDTH);
  int height = reader.get(cv::CAP_PROP_FRAME_HEIGHT);
  double fps = reader.get(cv::CAP_PROP_FPS);
  if (fps < 1 || fps > 50) { fps = 15.;}
  printf("infer_mp4 width: %d, height: %d, fps: %f\n", width, height, fps);

  // std::string writer_path = mp4_path + ".rec.mp4";
  // cv::VideoWriter writer = cv::VideoWriter(writer_path, cv::VideoWriter::fourcc('a','v','c','1'), fps, cv::Size(width, height), true);

  int frame_cnt = 0;
  int infer_id = 0;
  int infer_plate_id = 0;
  stream::ImageBlob car_blob(stream::ImageBlobMode_BGR);
  stream::ImageBlob plate_blob(stream::ImageBlobMode_BGR);
  stream::ImageBlob draw_data(stream::ImageBlobMode_BGR);
  while(true) {
    // if (frame_cnt > 500) {break;}
    
    cv::Mat img;
    bool ok = reader.read(img);
    if (!ok) {break;}
    // frame_cnt++;

    cv::Mat draw_img = img.clone();
    std::vector<cv::Rect> boxs;
    rec_car(img, boxs);

    for (auto& car_box : boxs) {
      cv::rectangle(draw_img, car_box, cv::Scalar(255,0,0), 3);
      // std::cout << "img:" << img.size() << ", " << car_box << std::endl;

      cv::Mat car_img = img(car_box).clone();
      car_blob.source_id = frame_cnt++;
      car_blob.img = car_img;
      car_blob.model_id = "plate_det";
      model->set_data(car_blob.model_id, infer_id, car_blob);
      if (infer_id < 0) {continue;}

      ai::InferResult infer_res;
      while ((ret = model->get_data(car_blob.model_id, infer_id, infer_res)) != 0) {
        std::this_thread::sleep_for(milliseconds(3));
      }
      // std::cout << "det_infos.size():" << det_infos.size() << std::endl;

      for (auto& det_info : infer_res.infer_res) {
        ai::DetInfo* di = (ai::DetInfo*)det_info;
        // std::cout << "car_img:" << car_img.size() << ", " << di->det_box << std::endl;

        // cv::rectangle(blob.img, di->det_box, cv::Scalar(0, 0, 255), 2);
        cv::Mat plate_img = car_img(di->det_box).clone();
        frame_cnt++;
        cv::imwrite(std::to_string(frame_cnt) + ".plate.jpg", plate_img);
        cv::imwrite(std::to_string(frame_cnt) + ".car.jpg", car_img);


        plate_blob.source_id = frame_cnt++;
        plate_blob.img = plate_img;
        plate_blob.model_id = "plate_rec";
        model->set_data(plate_blob.model_id, infer_plate_id, plate_blob);
        if (infer_plate_id < 0) {continue;}

        ai::InferResult plate_infer_res;
        while ((ret = model->get_data(plate_blob.model_id, infer_plate_id, plate_infer_res)) != 0) {
          std::this_thread::sleep_for(milliseconds(3));
        }

        if (plate_infer_res.infer_res.size() > 0) {
          ai::TextRecInfo* plate_rec_ifo = (ai::TextRecInfo*)plate_infer_res.infer_res[0];

          std::vector<stream::OSDDrawInfo> osdd_ifos;
          osdd_ifos.push_back(stream::OSDDrawInfo(di->det_box.x+car_box.x, di->det_box.y+car_box.y, di->det_box.width, di->det_box.height,di->score));
          osdd_ifos.back().texts.push_back("N"+std::to_string(di->score));
          osdd_ifos.back().texts.push_back("N"+plate_rec_ifo->text+":"+std::to_string(plate_rec_ifo->score));
          osdd_ifos.back().texts.push_back("N"+plate_rec_ifo->color);

          draw_data.img = draw_img;
          osd->draw_osd(draw_data, osdd_ifos);
          cv::imwrite(std::to_string(frame_cnt) + ".draw.jpg", draw_img);
          delete plate_infer_res.infer_res[0];
        }

        delete det_info;
      }

    }

    // writer.write(draw_img); 
    printf("frame_cnt:%d\n", frame_cnt);
  }
  reader.release();
  // writer.release();

}




int infer_det_img_plate(ai::RKModelManager* model, stream::CVOsdText* osd, std::string mp4_path) {
  int ret = 0;

  std::vector<boost::filesystem::path> all_files;
  get_files_in_dir_recursive(mp4_path, all_files, {".jpg", ".png"});
  // std::vector<std::pair<int64_t, std::string>> img_paths;
  // for (auto& f : all_files) { 
  //   std::string filename = f.filename().string();
  //   std::vector<std::string> ips;
  //   boost::algorithm::split(ips, filename, boost::is_any_of("."), boost::token_compress_on);
  //   int64_t time_point = std::stoll(ips[1]);
  //   img_paths.push_back(std::make_pair(time_point, f.string()));
  // }
  // sort(img_paths.begin(), img_paths.end(), [](std::pair<int64_t, std::string> a, std::pair<int64_t, std::string> b){return a.first < b.first;});

  int frame_cnt = 0;
  int infer_id = 0;
  int infer_plate_id = 0;
  stream::ImageBlob car_blob(stream::ImageBlobMode_BGR);
  stream::ImageBlob plate_blob(stream::ImageBlobMode_BGR);
  stream::ImageBlob draw_data(stream::ImageBlobMode_BGR);
    // if (frame_cnt > 500) {break;}
  for (auto& img_path : all_files) {
    

    cv::Mat car_img = cv::imread(img_path.string());
    cv::Mat draw_img = car_img.clone();

    car_blob.source_id = frame_cnt++;
    car_blob.img = car_img;
    car_blob.model_id = "plate_det";
    model->set_data(car_blob.model_id, infer_id, car_blob);
    if (infer_id < 0) {continue;}

    ai::InferResult infer_res;
    while ((ret = model->get_data(car_blob.model_id, infer_id, infer_res)) != 0) {
      std::this_thread::sleep_for(milliseconds(3));
    }
    // std::cout << "det_infos.size():" << det_infos.size() << std::endl;

    int plate_idx=0;
    for (auto& det_info : infer_res.infer_res) {
      ai::DetInfo* di = (ai::DetInfo*)det_info;
      // std::cout << "car_img:" << car_img.size() << ", " << di->det_box << std::endl;

      // cv::rectangle(blob.img, di->det_box, cv::Scalar(0, 0, 255), 2);
      cv::Mat plate_img = car_img(di->det_box).clone();
      frame_cnt++;
      // cv::imwrite(std::to_string(frame_cnt) + ".plate.jpg", plate_img);
      // cv::imwrite(std::to_string(frame_cnt) + ".car.jpg", car_img);


      plate_blob.img = plate_img;
      plate_blob.model_id = "plate_rec";
      model->set_data(plate_blob.model_id, infer_plate_id, plate_blob);
      if (infer_plate_id < 0) {continue;}

      ai::InferResult plate_infer_res;
      while ((ret = model->get_data(plate_blob.model_id, infer_plate_id, plate_infer_res)) != 0) {
        std::this_thread::sleep_for(milliseconds(3));
      }

      if (plate_infer_res.infer_res.size() > 0) {
        ai::TextRecInfo* plate_rec_ifo = (ai::TextRecInfo*)plate_infer_res.infer_res[0];

        std::vector<stream::OSDDrawInfo> osdd_ifos;
        osdd_ifos.push_back(stream::OSDDrawInfo(di->det_box.x, di->det_box.y, di->det_box.width, di->det_box.height,di->score));
        osdd_ifos.back().texts.push_back("N"+std::to_string(di->score));
        osdd_ifos.back().texts.push_back("N"+plate_rec_ifo->text+":"+std::to_string(plate_rec_ifo->score));
        osdd_ifos.back().texts.push_back("N"+plate_rec_ifo->color);
        std::cout << img_path.string() << ", " << plate_rec_ifo->text << std::endl;
        draw_data.img = draw_img;
        osd->draw_osd(draw_data, osdd_ifos);
        cv::imwrite(img_path.string() +"." +std::to_string(plate_idx++) + ".draw.jpg", draw_img);
        delete plate_infer_res.infer_res[0];
      }

      delete di;

    }

    // writer.write(draw_img); 
    // printf("frame_cnt:%d\n", frame_cnt);
  }
}



int main(int argc, char *argv[]){
  std::string model_dir = argv[1];

  int ret = 0;
  LogInfo *log_info = new LogInfo(1, "log", "log");

  // 替代 opencv putText 函数 画中文
  stream::CVStreamOsdInfo osd_ifo;
  osd_ifo.boxColorOK = cv::Scalar(255,0,0);
  osd_ifo.boxColorNG = cv::Scalar(0,0,255);
  osd_ifo.fontHeight = 48;
  stream::CVOsdText* osd = new stream::CVOsdText(&osd_ifo, log_info);
  if ((ret = osd->init_params()) != 0) {
    std::cout << "osd->init_params error: " << ret << std::endl;
    return ret;
  }

  ai::RKModelManager* model = new ai::RKModelManager(log_info);

  // // 测试加密模型时需要
  // //初始化授权
  // ret = model->init_param(log_info, "Rockchip");
  // if(ret != 0){
  //   std::cout << "ERROR. init_param ret:" << ret << std::endl; 
  //   return ret;
  // }
  // //获取所支持的算法标签
  // std::vector<au::authModelLabel> auth_mdls;
  // ret = model->get_all_model_info(auth_mdls);


  std::vector<ai::InitModelData> init_model_datas;
  // //加密模型初始化示例
  // init_model_datas.push_back(ai::InitModelData(model_dir, "", "det_chongdianzhan_yolom.20241231_001", "", ai::model_det, 0, "det_chongdianzhan_yolom.20241231_001.yaml"));
  // init_model_datas.push_back(ai::InitModelData(model_dir, "", "det_chongdianqiang_motuoche_yolos.20250210", "", ai::model_det, 1, "det_chongdianqiang_motuoche_yolos.20250210.yaml"));
  //原始模型初始化示例
  // init_model_datas.push_back(ai::InitModelData(model_dir, "cls/firesmoke", "xfc3s32d", "firesmoke", ai::model_cls_video, 0, "cfg.yaml"));
  // init_model_datas.push_back(ai::InitModelData(model_dir, "det/common5", "xfc3s32d", "common5", ai::model_det, 0, "cfg.yaml"));
  // init_model_datas.push_back(ai::InitModelData(model_dir, "ocr/plate/det", "23sdxfc3", "plate_det", ai::model_det, 0, "cfg.yaml"));
  // init_model_datas.push_back(ai::InitModelData(model_dir, "ocr/plate/rec", "34sdf34z", "plate_rec", ai::model_text_rec, 1, "cfg.yaml"));
  // init_model_datas.push_back(ai::InitModelData(model_dir,"det/facedet", "67fgwssdh", "facedet", ai::model_face_det, 0, "cfg.yaml"));
  // init_model_datas.push_back(ai::InitModelData(model_dir,"det/facerec", "sdh67fgws", "facerec", ai::model_face_rec, 1, "cfg.yaml"));

  ret = model->init_model(init_model_datas, 0);
  if (ret != 0) { 
    std::cout << "init_model ret: " << ret << std::endl;
    return ret; 
  }

  std::this_thread::sleep_for(std::chrono::milliseconds(3000));
  
  std::string file_path = argv[2];
  // // 性能测试 支持多模型
  // int max_size = 4;
  // performance_test(model, file_path, init_model_datas, max_size);

  // 图片 视频 识别
  // infer_det_mp4(model, osd, init_model_datas, file_path);
  // infer_det_mp4_plate(model, osd, file_path);
  // infer_det_img_plate(model, osd, file_path);
  infer_det_img(model, osd, init_model_datas, file_path);
  // infer_cls_video(model, osd, init_model_datas, file_path);
  // infer_det_img_to_via_file(model, osd, init_model_datas, file_path);
  // infer_face_img_det_rec(model, osd, file_path);
  // 多模型同时推理
  // infer_det_imgmulti(model, osd, init_model_datas, file_path);



  // // 启动http服务端
  // Server svr;
  // svr.Get("/hi", [](const Request& req, Response& res) {
  //   res.set_content("c++ service for zqdl smartbox start ok!", "text/plain");
  // });
  // svr.Post("/api/rec", [&](const Request& req, Response& res) {
  //   std::string return_msg;
  //   int ret = 0;
  //   if (0 != (ret = infer_det_img_api(model, init_model_datas, req.body, return_msg))) {
  //     return_error(ret, MyError::GetErrorString(ret), return_msg);
  //   }
  //   res.set_content(return_msg, "application/json");
  // });
  // svr.listen("0.0.0.0", 9918);
  return 0;
}
